
import matplotlib.pyplot as plt
import numpy as np
if __name__ == "__main__":

    fs = 5000
    t_max = 1
    n = int(fs*t_max)

    # Build signal
    t = np.linspace(0,t_max,n)
    dt = t[1]-t[0]


    # Define a time-dependent phase function
    f_base = 350
    df_1 = 10
    tau_1 = 0.5
    omega_1 = lambda x: f_base + (df_1)*(1-np.exp(-x/tau_1))
    amp_1 = lambda x: 1 + 0.5*np.exp(-x/3)


    omega_2 = lambda x: 80-2*x
    amp_2 = lambda x: 8 - 0.5*np.exp(-x)


    mu = lambda x: 1.5 + 2.5*np.exp(-x/(1.5))
    #mu = lambda x: 2*x-0.3*x**2
    #mu = lambda t: (0.2)*np.sin(t)

    # Create data with time-dependent frequency
    z_1 = lambda t: amp_1(t)*np.cos(2*np.pi*omega_1(t)*t)
    z_2 = lambda t: amp_2(t)*np.cos(2*np.pi*omega_2(t)*t)

    z = z_1(t)
    z += z_2(t)
    z += mu(t)
    #z = z/np.std(z)

    stft = plt.specgram(z, NFFT=4800, Fs=fs, noverlap=256)
    plt.figure()
    psd = plt.psd(z, NFFT=1200, Fs=fs)
    plt.show()